#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : wes_check_cnv.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/04 09:56
@Description :
    检查 CNV 结果文件中是否存在与目标基因区间重叠的 CNV 片段。
    支持多个样本输入与批量结果汇总。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import pandas as pd


def wes_check_cnv(args: Namespace):
    """
    检查 CNV 片段与目标基因区间的重叠情况。

    参数
    ----
    args : Namespace
        命令行参数对象，包含以下属性：
            - chr_name (str): 染色体名称，例如 "chr5"
            - gene_start (int): 基因起始位置
            - gene_end (int): 基因终止位置
            - samples (str): 样本名称，逗号分隔
            - in_files (str): CNV 文件路径，逗号分隔，与 samples 对应
            - out_file (Path): 输出 Excel 文件路径
    """

    chr_name = args.chr_name.strip()
    gene_start = args.gene_start
    gene_end = args.gene_end

    # 解析样本名和输入文件
    samples = [s.strip() for s in args.samples.strip().split(",")]
    in_files = [Path(i).resolve() for i in args.in_files.strip().split(",")]

    # 创建输出目录
    out_file = args.out_file.resolve()
    out_file.parent.mkdir(parents=True, exist_ok=True)

    out_dfs = []

    # 遍历每个样本
    for sample, in_file in zip(samples, in_files):
        # 读取 CNV 文件
        df = pd.read_csv(in_file, sep="\t", encoding="utf-8")

        # 添加样本名列
        df["Sample"] = sample

        # 判断 CNV 片段是否与目标区间重叠
        overlap_mask = (
            (df["chromosome"] == chr_name)
            & (df["start"] <= gene_end)
            & (df["end"] >= gene_start)
        )

        # 筛选重叠的 CNV 片段
        overlap_df = df.loc[overlap_mask, :]

        if not overlap_df.empty:
            out_dfs.append(overlap_df)

    # 合并所有样本结果
    if out_dfs:
        result_df = pd.concat(out_dfs, axis=0, ignore_index=True)
    else:
        # 若无重叠结果，输出空表结构
        result_df = pd.DataFrame(columns=["chromosome", "start", "end", "Sample"])

    # 导出到 Excel 文件
    result_df.to_excel(out_file, index=False, engine="openpyxl")


def main():
    """命令行入口"""
    parser = ArgumentParser(description="检查 CNV 片段与目标基因的重叠情况")
    parser.add_argument("--samples", type=str, required=True, help="样本名称，逗号分隔")
    parser.add_argument("--in_files", type=str, required=True, help="CNV 文件路径，逗号分隔")
    parser.add_argument("--out_file", type=Path, required=True, help="输出 Excel 文件路径")
    parser.add_argument("--chr_name", type=str, default="chr5", help="染色体名称")
    parser.add_argument("--gene_start", type=int, default=10564442, help="基因起始位置")
    parser.add_argument("--gene_end", type=int, default=10650308, help="基因终止位置")
    args = parser.parse_args()

    wes_check_cnv(args)


if __name__ == "__main__":
    main()
